from asyncore import read
from copy import copy
from datetime import datetime
from io import StringIO
import json
import os
import sys
import time
import pymongo

from data_reader.DataReader import DataReader
from data_reader.NsReader import NsReader
from data_reader.OrReader import OrReader
from utils.Timer import Timer

def connect():
    conn = pymongo.MongoClient('127.0.0.1', 27017)
    mongo_db = conn['picsad']
    return mongo_db['point_cloud']

def query_circle(table, x, y, radius):
    timer = Timer()
    timer.Start()
    cur = connect()
    res_f = open("result.txt", "w")
    points = cur.find({'x': {'$gt': x-radius}, 'x': {'$lt': x+radius}, 
              'y': {'$gt': y-radius}, 'y': {'$lt': y+radius}})
    for p in points:
        x_s = p['x'] - x
        y_s = p['y'] - y
        if x_s*x_s < radius*radius and y_s*y_s < radius*radius:
            res_f.write(str(p['x']))
            res_f.write(str(p['y']))
            res_f.write(str(p['x']))

    res_f.close()
    timer.StopAndRecord('query circle %d %d %d'%(x, y, radius))
    timer.PrintAll()
    return len(list(points))

def query_rect(table, x_min, x_max, y_min, y_max):
    cur = connect()
    timer = Timer()
    timer.Start()
    cur = connect()
    res_f = open("result.txt", "w")
    points = cur.find({'x': {'$gt': x_min}, 'x': {'$lt': x_max}, 
              'y': {'$gt': y_min}, 'y': {'$lt': y_max}})
    for p in points:
        res_f.write(str(p['x']))
        res_f.write(str(p['y']))
        res_f.write(str(p['z']))
    res_f.close()
    timer.StopAndRecord('query rect')
    timer.PrintAll()
    return len(list(points))

def query_sweep(table, sweep_index):
    cur = connect()
    timer = Timer()
    timer.Start()
    cur = connect()
    res_f = open("result.txt", "w")
    points = cur.find({'sweep': sweep_index})
    for p in points:
        res_f.write(str(p['x']))
        res_f.write(str(p['y']))
        res_f.write(str(p['x']))
    res_f.close()
    timer.StopAndRecord('query sweep')
    timer.PrintAll()
    return len(list(points))

def query_timestamp(table, tstamp_start, tstamp_end):
    cur = connect()
    timer = Timer()
    timer.Start()
    cur = connect()
    res_f = open("result.txt", "w")
    points = cur.find({'tstamp': {'$gt': tstamp_start}, 'tstamp': {'$lt': tstamp_end}})
    for p in points:
        res_f.write(str(p['x']))
        res_f.write(str(p['y']))
        res_f.write(str(p['x']))
    res_f.close()
    timer.StopAndRecord('query timestamp')
    timer.PrintAll()
    return len(list(points))


def query_classification(table, clss_arr):
    cur = connect()
    timer = Timer()
    timer.Start()
    cur = connect()
    res_f = open("result.txt", "w")
    for clss in clss_arr:
        points = cur.find({'cls': clss})
        for p in points:
            res_f.write(str(p['x']))
            res_f.write(str(p['y']))
            res_f.write(str(p['x']))
    res_f.close()
    timer.StopAndRecord('query clas')
    timer.PrintAll()
    return len(list(points))


if __name__ == '__main__':
    query_circle("point_cloud", 50, 100, 20)
    query_circle("point_cloud", 0, 50, 10)
    query_rect("point_cloud", 40, 70, 200, 240)
    query_rect("point_cloud", -20, 60, 20, 100)
    query_sweep("point_cloud", 1)
    query_sweep("point_cloud", 10)
    query_timestamp("point_cloud", 1317617734.934675, 1317617738.252331)
    query_classification("point_cloud", [10])
    query_classification("point_cloud", [80, 81, 99])